#!/usr/bin/env python
#-*-coding:utf-8-*-

import pre_process

from predict import Predictor
import conf

predictor = Predictor(map_file=conf.SEG_MAP_PATH, checkpoint_dir=conf.SEG_CHECKPOINT_DIR)


def seg_by_local_model(text):
    text = pre_process.pre_handle_text(text)
    words = []
    seg_result = []

    if not text:
        return seg_result, []

    labels = predictor.predict(text)
    N = len(labels)
    tmp_word = ""

    seg_result = [[w, l] for w, l in zip(list(text), labels)]

    for i in range(N):
        label = labels[i]
        w = text[i]
        if label == "B":
            if tmp_word:
                words.append(tmp_word)
                tmp_word = ""
            tmp_word += w
        elif label == "M":
            tmp_word += w
        elif label == "E":
            tmp_word += w
            words.append(tmp_word)
            tmp_word = ""
        else :
            tmp_word = ""
            words.append(w)
    if tmp_word:
        words.append(tmp_word)
    return seg_result, words



if __name__ == '__main__':

    text = "刚刚在英国完成大婚的周杰伦，又有新消息传来。不过，不是粉丝们翘首以盼的“昆凌有喜”的信息，而是第四季《中国好声音》公布首位导师，周董确认加盟，成为“好声音”史上第八位导师。"
    seg_result , words = seg_by_local_model(text)
    print(words)
